import os
import sys
import time
import cv2
import numpy as np
import torch
import torch.optim as optim
import torch.backends.cudnn as cudnn
from torch.utils.data import DataLoader

from models.efficientdet import EfficientDet
from datasets.defect.insdefectloader import DefectAnno
from utils import EFFICIENTDET


def inferenceBatch(model, batIm, dim, rawIm, label=None, isShow=False):
    for i, img in enumerate(batIm):
        inference(model, img, dim, rawIm[i], label=label[i], isShow=isShow)


def inference(model, origin_img, dim, rawIm, label=None, isShow=False):
    model.eval()
    model.is_training = False
    scores, classification, transformed_anchors = model(torch.unsqueeze(origin_img, 0))
    bboxes = list()

    if isShow:
        showIm = rawIm.numpy()
        bbox = torch.squeeze(label, 0).detach().cpu().numpy().astype(np.int)
        showIm = cv2.rectangle(showIm, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])),
                               (255, 0, 255), 2)
        cls = bbox[4]
        cv2.putText(
            showIm, '{} {}'.format(cls, int(100)),
            (bbox[0], bbox[1]), cv2.FONT_HERSHEY_SIMPLEX,
            0.8, (255, 0, 255), 2
        )

        labelIm = cv2.imread("labelName.png")
        cv2.imshow("label", labelIm)

    for j in range(scores.shape[0]):
        if scores[j] > 0.5 and classification[j] == cls:
            bbox = transformed_anchors[[j], :][0].data.cpu().numpy()
            x1 = int(bbox[0])
            y1 = int(bbox[1])
            x2 = int(bbox[2])
            y2 = int(bbox[3])
            bboxes.append([x1, y1, x2, y2])
            if isShow:
                cv2.rectangle(showIm, (x1, y1),
                              (x2, y2), (255, 0, 0), 2)
                cv2.putText(
                    showIm, '{} {}'.format(classification[j], int(scores[j])),
                    (x1, y1), cv2.FONT_HERSHEY_SIMPLEX,
                    0.8, (255, 0, 0), 2
                )

    if isShow:
        cv2.imshow("testOut", showIm)
        cv2.waitKey(25)

    model.train()
    model.is_training = True
